package com.wish233.register.impl;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.json.JSONUtil;
import com.wish233.config.RegistryConfig;
import com.wish233.model.register.ServiceMetaInfo;
import com.wish233.register.Registry;
import com.wish233.utils.HeartBeatUtils;
import com.wish233.utils.ServiceMetaUtils;
import io.etcd.jetcd.*;
import io.etcd.jetcd.options.GetOption;
import io.etcd.jetcd.options.PutOption;
import io.etcd.jetcd.watch.WatchEvent;
import io.vertx.core.impl.ConcurrentHashSet;
import lombok.extern.slf4j.Slf4j;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

/**
 * @author WHH
 */
@Slf4j
public class EtcdRegistry implements Registry {

    public static final String ETCD_ROOT_PATH = "/wrpc/";
    private final Set<String> watchingKeySet = new ConcurrentHashSet<>();
    private Client client;
    private KV kvClient;
    private RegistryServiceCache cache;

    @Override
    public void init(RegistryConfig registryConfig) {
        //初始化客户端
        client = Client.builder().endpoints(registryConfig.getAddress())
                .connectTimeout(Duration.ofMillis(registryConfig.getTimeout()))
                .build();
        kvClient = client.getKVClient();
        heartBeat();
        Runtime.getRuntime().addShutdownHook(new Thread(this::destroy));
        cache = RegistryServiceCache.getCache();
    }

    @Override
    public void register(ServiceMetaInfo serviceMetaInfo) throws Exception {
        String serviceNodeKey = serviceMetaInfo.getServiceNodeKey();
        Lease leaseClient = client.getLeaseClient();
        //添加一个租约 设置为30秒
        long leaseId = leaseClient.grant(30).get().getID();
        String regKey = ETCD_ROOT_PATH + serviceNodeKey;
        ByteSequence key = ByteSequence.from(regKey, StandardCharsets.UTF_8);
        ByteSequence value = ByteSequence.from(JSONUtil.toJsonStr(serviceMetaInfo), StandardCharsets.UTF_8);
        kvClient.put(key, value);
        PutOption putOption = PutOption.builder().withLeaseId(leaseId).build();
        kvClient.put(key, value, putOption);
        //添加节点到heartBeatSet
        HeartBeatUtils.addNode(regKey);
    }

    @Override
    public void unRegister(ServiceMetaInfo serviceMetaInfo) {
        String regKey = ETCD_ROOT_PATH + serviceMetaInfo.getServiceKey();
        ByteSequence key = ByteSequence.from(regKey, StandardCharsets.UTF_8);
        kvClient.delete(key);
        //删除心跳节点
        HeartBeatUtils.removeNode(regKey);
    }

    @Override
    public List<ServiceMetaInfo> serviceDiscovery(String serviceKey) {
        //读缓存
        List<ServiceMetaInfo> serviceMetaInfos = cache.readCache(serviceKey);
        if (CollectionUtil.isNotEmpty(serviceMetaInfos)) {
            return serviceMetaInfos;
        }
        //查询一个目录下所有的服务
        String prefix = ETCD_ROOT_PATH + serviceKey + "/";
        GetOption getOption = GetOption.builder().isPrefix(true).build();
        try {
            List<KeyValue> kvs = kvClient.get(
                    ByteSequence.from(prefix, StandardCharsets.UTF_8),
                    getOption).get().getKvs();
            //写缓存
            List<ServiceMetaInfo> metaInfoList = kvs.stream().map(kv -> {
                String key = kv.getKey().toString(StandardCharsets.UTF_8);
                String value = kv.getValue().toString(StandardCharsets.UTF_8);
                //监听缓存
                watch(key);
                return JSONUtil.toBean(value, ServiceMetaInfo.class);
            }).collect(Collectors.toList());
            //监听
            cache.writeCache(serviceKey, metaInfoList);
            return metaInfoList;
        } catch (InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }

    }

    @Override
    public void heartBeat() {
        HeartBeatUtils.doHeartBeat(key -> {
            try {
                List<KeyValue> kvs = kvClient.get(ByteSequence.from(key, StandardCharsets.UTF_8)).get().getKvs();
                if (CollectionUtil.isEmpty(kvs)) {
                    return;
                }
                //重新注册（续期）
                KeyValue kv = kvs.get(0);
                String value = kv.getValue().toString(StandardCharsets.UTF_8);
                ServiceMetaInfo serviceMetaInfo = JSONUtil.toBean(value, ServiceMetaInfo.class);
                register(serviceMetaInfo);
            } catch (Exception e) {
                throw new RuntimeException(key + "续期失败", e);
            }
        });
    }

    @Override
    public void destroy() {
        Set<String> set = HeartBeatUtils.getSet();
        //下线所有节点
        for (String key : set) {
            try {
                kvClient.delete(ByteSequence.from(key, StandardCharsets.UTF_8)).get();
            } catch (InterruptedException | ExecutionException e) {
                throw new RuntimeException(key + "节点下线失败", e);
            }
        }
        log.info("节点下线！");
        Optional.of(client).ifPresent(Client::close);
        Optional.of(kvClient).ifPresent(KV::close);
    }

    /**
     * 监听（消费端）
     *
     * @param serviceNodeKey
     */
    @Override
    public void watch(String serviceNodeKey) {
        Watch watchClient = client.getWatchClient();
        // 之前未被监听，开启监听
        boolean newWatch = watchingKeySet.add(serviceNodeKey);
        if (newWatch) {
            watchClient.watch(ByteSequence.from(serviceNodeKey, StandardCharsets.UTF_8), response -> {
                for (WatchEvent event : response.getEvents()) {
                    switch (event.getEventType()) {
                        // key 删除时触发
                        case DELETE:
                            // 清理注册服务缓存
                            String serviceKey = ServiceMetaUtils.extractServiceKeyFromNodeKey(serviceNodeKey, '/');
                            cache.clearCache(serviceKey);
                            break;
                        case PUT:
                        default:
                            break;
                    }
                }
            });
        }
    }

}
